import numpy as np
import torch
from model import Vit
from dataset import Flowers
import matplotlib.pyplot as plt
from torchvision import transforms
from PIL import Image

if __name__ == '__main__':
    # 建立数据集
    data_transform = transforms.Compose([transforms.Resize(256),
                                         transforms.CenterCrop(224),
                                         transforms.ToTensor(),
                                         transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
    img = Image.open('datasets/flower_data/val/roses/505517255_cfbb6f6394.jpg')
    input = data_transform(img).unsqueeze(0)
    label2class = Flowers(dataset_path='datasets/flower_data/val').label2class

    device = torch.device('cuda:0')

    # 建立模型
    model = Vit(img_size=[224, 224],
                patch_size=16,
                embed_dim=192,
                depth=6,
                num_heads=4,
                num_classes=5).to(device)
    weight_dict = torch.load('vit.pth', map_location=device)
    model.load_state_dict(weight_dict)

    model.eval()
    with torch.no_grad():
        output = model(input.to(device))
        output = output.detach().cpu()

        label = output[0].numpy().argmax()
        cnf = torch.softmax(output[0],dim=0).numpy().max()*100.0
        cnf = np.around(cnf, decimals=2) #保留2位小数
    plt.imshow(img)
    plt.title('{} : {}%'.format(label2class[label],cnf))
    plt.show()
